from opt.treesearch import TreeNode as BaseTreeNode, DFS as BaseDFS

from utils.multiset import MultiSet
from utils.namespace import Namespace
from utils.misc import INF

from cranesched.solution import Solution


class TreeNode(BaseTreeNode):
    """Tree node class for the Crane Scheduling problem for non-productive minimization. 
    Implements the basic operations required by the DFS tree search solver."""
    def objective(self):
        return (self.solution.total_penalty, self.solution.total_traveled_distance)
        
    def lower_bound_with(self, assignment):
        """Compute the lower bound of the solution obtained by appending 'assignment' to our 
        current solution."""
        solution = self.solution
        return (solution.total_penalty + assignment.penalty, 
                solution.total_traveled_distance + assignment.traveled_distance)
        
    def solution_and_meta(self):
        """Return the solution data and metadata contained in the tree node."""
        return [a.id for a in self.assignments], self.count
        
    @classmethod
    def root(cls, solver):
        self = cls()
        self.solver = solver
        self.solution = solver.global_solution
        self.assignments = []
        self.restricted_to = solver.restricted_to
        self.count = MultiSet()
        return self
        
    def branches(self):
        # build a list of assignments for our assignable activities
        solution = self.solution
        assignments = []
        assignable = (solution.activities.assignable if self.restricted_to is None else 
                      solution.activities.assignable & self.restricted_to)
        for activity in assignable:
            assignments.extend(solution.eligible_assignments(activity))
        if len(assignments) == 0:
            self.solver.log.debug("infeasible solution")
        
        # compute lower bound for each branch, prune non-improving branches, and sort the remaining
        # branches by decreasing order of lower bound. The order is decreasing because the branches
        # are reversed when being pushed into the solver's stack, i.e. the last branch goes to the 
        # top of the stack, therefore it is the first to be explored.
        LB = dict()
        for assignment in assignments:
            lower_bound = self.lower_bound_with(assignment)
            if lower_bound < self.solver.upper_bound:
                LB[assignment] = (lower_bound, assignment.id)
            else:
                self.count.add("pruned")
        return sorted(LB.iterkeys(), key=LB.__getitem__, reverse=True)
        
    def apply(self, assignment):
        self.solution.append(assignment)
        self.count.add("descents")
        
    def undo(self, assignment):
        assert self.solution.pop() is assignment
        self.count.add("backtracks")
        
    def is_leaf(self):
        is_leaf = self.solution.is_complete
        if is_leaf:
            self.count.add("leaves")
        return is_leaf
        
        
class DFS(BaseDFS):
    initial_upper_bound = (INF, INF)
    initial_lower_bound = (0.0, 0.0)
    default_pruning = False  # pruning is done in the branches() method
    default_root_fnc = staticmethod(TreeNode.root)
    
    @BaseDFS.initialized
    def initialized(self, instance, global_solution=None, subinstance_size=None, *args, **kwargs):
        BaseDFS.initialized.enter(self, instance, *args, **kwargs)
        if global_solution is None:
            self.global_solution = Solution(self.instance)
        else:
            self.global_solution = global_solution
        self.restricted_to = (None if subinstance_size is None else 
                              self.global_solution.activities.feasible_subset(subinstance_size))
        
